#!/bin/bash
###############################################################
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# installer is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
###############################################################

DRIVER_INSTALL_EXEC_DIR=$1
OPERATE_TYPE=$2

function _log() {
	local prefix="$1"
	shift
	echo "$(date +"[%Y-%m-%d %H:%M:%S,%N]") [${prefix}] $*"
}

function info_log () {
    _log "INFO" "$*"
}

function warning_log () {
    _log "WARNING" "$*"
}

function error_log () {
    _log "ERROR" "$*"
}

function fatal_log () {
    _log "FATAL" "$*"
    exit 1
}

function install_basic_packages() {
    info_log "Install dependency packages"

    if command -v yum >/dev/null 2>&1; then
        info_log "Using yum for package installation"
        yum clean all
        yum makecache
        PACKAGE_MANAGER="yum"
        PKG_LIST="jq wget unzip which initscripts coreutils findutils gawk e2fsprogs util-linux net-tools pciutils gcc make automake autoconf libtool git patch kernel-devel-$(uname -r) kernel-headers-$(uname -r) dkms"

    elif command -v apt-get >/dev/null 2>&1; then
        info_log "Using apt-get for package installation"
        apt-get update
        PACKAGE_MANAGER="apt-get"
        PKG_LIST="jq wget unzip debianutils coreutils findutils gawk e2fsprogs util-linux net-tools pciutils gcc make automake autoconf libtool git patch dkms linux-headers-$(uname -r)"

    elif command -v dnf >/dev/null 2>&1; then
        info_log "Using dnf for package installation"
        dnf makecache
        PACKAGE_MANAGER="dnf"
        PKG_LIST="jq wget unzip which initscripts coreutils findutils gawk e2fsprogs util-linux net-tools pciutils gcc make automake autoconf libtool git patch kernel-devel-$(uname -r) kernel-headers-$(uname -r) dkms"

    else
        fatal_log "No compatible package manager found (yum, apt-get, dnf)."
    fi

    for pkg in $PKG_LIST; do
        case "$PACKAGE_MANAGER" in
            yum|dnf)
                if ! rpm -q "$pkg" >/dev/null 2>&1; then
                    info_log "Installing $pkg..."
                    if ! $PACKAGE_MANAGER -y install "$pkg"; then
                        fatal_log "Failed to install $pkg using $PACKAGE_MANAGER"
                    fi
                else
                    info_log "$pkg is already installed, skipping..."
                fi
                ;;
            apt-get)
                if ! dpkg -l | grep -w "^ii  $pkg"; then
                    info_log "Installing $pkg..."
                    if ! $PACKAGE_MANAGER -y install "$pkg"; then
                        fatal_log "Failed to install $pkg using $PACKAGE_MANAGER"
                    fi
                else
                    info_log "$pkg is already installed, skipping..."
                fi
                ;;
        esac
    done

    info_log "Dependency packages installed successfully."
}

function get_chip_type() {
    CHIP_TYPE=$(lspci | grep -i d802)
    if [ -n "$CHIP_TYPE" ]; then
        CHIP_TYPE="910b"
    else
        CHIP_TYPE=$(lspci | grep -i d500)
        if [ -n "$CHIP_TYPE" ]; then
            CHIP_TYPE="310p"
        else
            fatal_log "Unsupported chip type"
        fi
    fi
    echo $CHIP_TYPE
}

function get_node_architecture() {
    NODE_ARCH=$(uname -m)
    case $NODE_ARCH in
        "aarch64"|"x86_64")
            echo $NODE_ARCH
            ;;
        *)
            fatal_log "Unsupported architecture: $NODE_ARCH"
            ;;
    esac
}

function get_version() {
  json_output=$(kubectl get npuclusterpolicies.npu.openfuyao.com -o json)


  if [ -z "$json_output" ]; then
    fatal_log "No output from kubectl command."
  fi

  output=$(echo "$json_output" | jq -r '.items[0].spec.driver.version')

  if [ "$output" == "null" ]; then
    fatal_log "spec.driver.version is null or missing."
  fi

  echo "$output"
}

function download_and_extract() {
  VERSION=$(get_version)
  json_file="/opt/${DRIVER_INSTALL_EXEC_DIR}/NPU/$VERSION/config.json"
  echo "File: $json_file"

  chip_type=$(get_chip_type)
  arch=$(get_node_architecture)
  chip_type_arch="${chip_type}_${arch}"

  json_value=$(jq -r ".other[] | select(has(\"$chip_type_arch\")) | .\"$chip_type_arch\"" "$json_file")

  if [ -z "$json_value" ]; then
    fatal_log "Error: No matching entry found for key '$chip_type_arch' in the configuration file."
  else
    echo "Found value for $chip_type_arch: $json_value"
  fi

  

  download_url=$(echo "$json_value" | jq -r '.url')

  install_path="/opt/${DRIVER_INSTALL_EXEC_DIR}/npu-driver/${chip_type}"
  main_zip="$install_path/$(basename "$download_url")"
  http_referer="https://www.hiascend.com/"
# 离线下载zip包放在路径： /opt/openFuyao-npu-driver/npu-driver/910b/ 下  
  if [ -f "$main_zip" ]; then
    echo "Local zip file found: $main_zip, skipping download"
  else
    mkdir -p "$install_path"
    echo "Local zip file not found: $main_zip, start to dowanload"
    if ! wget --referer=$http_referer -O "$install_path/$(basename "$download_url")" "$download_url"; then
        fatal_log "Error: Failed to download file from $download_url."
    fi
    echo "Download completed successfully."
  fi

  
  unzip -o "$main_zip" -d "/opt/${DRIVER_INSTALL_EXEC_DIR}/npu-driver"

  nested_zip=$(find "/opt/${DRIVER_INSTALL_EXEC_DIR}/npu-driver" -name "*.zip" -not -path "$main_zip" | head -n 1)
  if [ -n "$nested_zip" ]; then
    echo "Found nested ZIP file: $nested_zip"

    unzip -o "$nested_zip" -d "$install_path"

    find "/opt/${DRIVER_INSTALL_EXEC_DIR}/npu-driver" -mindepth 1 -not -path "/opt/${DRIVER_INSTALL_EXEC_DIR}/npu-driver/${chip_type}*" -exec rm -rf {} +

    mv "$install_path"/* /opt/${DRIVER_INSTALL_EXEC_DIR}/npu-driver/

    rm -rf "$install_path"
  fi
}


function exec_install_npu_driver () {
    install_basic_packages

    download_and_extract

    cd /opt/$DRIVER_INSTALL_EXEC_DIR/npu-driver || fatal_log "Failed to change directory to /opt/$DRIVER_INSTALL_EXEC_DIR/npu-driver"

    bash install.sh install all
    if [ $? -ne 0 ]; then
        fatal_log "Failed to install npu driver"
    fi
    info_log "success install npu driver"
}

if [ "$OPERATE_TYPE" = "install" ]; then
    exec_install_npu_driver
else
    error_log "Invalid operate type"
fi
